import os
import numpy as np
import nibabel as nib
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import cv2
import matplotlib.pyplot as plt
import torchvision.transforms as transforms

class MRIDataset(Dataset):
    def __init__(self, root_dir, modalities=('axi', 'sag', 'cor'), max_slices=3, 
                 crop_size=224, target_size=224, augment=None,transform=None):
        self.root_dir = root_dir
        self.modalities = modalities
        self.max_slices = max_slices
        self.crop_size = crop_size
        self.augment = augment
        self.target_size = target_size
        self.transform = transform
        # Load labels
        self.labels = pd.read_csv(os.path.join(root_dir, 'label.csv'),
                         header=0,
                         names=['p_id', 'label']).set_index('p_id')
        self.labels['label'] = self.labels['label'].astype(int) - 1  # 转换为0-based
        assert self.labels['label'].between(0, 2).all(), "标签必须在0-2范围内"
        
        self.patient_ids = []
        for pid in self.labels.index:
            if self._check_data_complete(pid):
                self.patient_ids.append(pid)

        self.class_counts = {i: 0 for i in range(2)}  # 0-based classes 0,1,2
        for pid in self.patient_ids:
            label = self.labels.loc[pid, 'label']
            self.class_counts[label] += 1

    def _check_data_complete(self, pid):
        # 与原代码相同，保持不变
        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            if not (os.path.exists(img_path) and os.path.exists(roi_path)):
                return False
            
            try:
                img = nib.load(img_path)
                roi = nib.load(roi_path)
            except Exception as e:
                print(f"加载文件失败：{img_path} 或 {roi_path}, 错误：{e}")
                return False
            
            img_shape = img.shape[:2]
            roi_shape = roi.shape[:2]
            
            if img_shape != roi_shape:
                print(f"患者 {pid} 的 {mod} 模态图像尺寸 {img_shape} 与 ROI 尺寸 {roi_shape} 不匹配")
                return False
            
            max_allowed_ratio = 1.2
            h, w = img_shape
            if h == 0 or w == 0:
                print(f"患者 {pid} 的 {mod} 模态存在零尺寸：{img_shape}")
                return False
            current_ratio = max(h / w, w / h)
            if current_ratio > max_allowed_ratio:
                print(f"患者 {pid} 的 {mod} 模态长宽比过大：{img_shape} (比例 {current_ratio:.2f})")
                return False
            
        return True

    def _crop_around_center(self, image, center_y, center_x, is_roi=False):
        """根据给定中心坐标进行裁剪"""
        target_size = self.crop_size
        half_size = target_size // 2
        
        # 计算理论裁剪范围
        y_start = center_y - half_size
        y_end = center_y + half_size
        x_start = center_x - half_size
        x_end = center_x + half_size
        
        # 计算实际可裁剪范围
        actual_y_start = max(0, y_start)
        actual_y_end = min(image.shape[0], y_end)
        actual_x_start = max(0, x_start)
        actual_x_end = min(image.shape[1], x_end)
        
        # 执行实际裁剪
        cropped = image[actual_y_start:actual_y_end,
                        actual_x_start:actual_x_end]
        
        # 计算需要填充的边界
        pad_y_before = max(0, -y_start)
        pad_y_after = max(0, y_end - image.shape[0])
        pad_x_before = max(0, -x_start)
        pad_x_after = max(0, x_end - image.shape[1])
        
        # 使用边缘像素填充
        padded = np.pad(cropped,
                        ((pad_y_before, pad_y_after),
                         (pad_x_before, pad_x_after)),
                        mode='edge')
        # img_resize = skimage.transform.resize(padded, (self.target_size, self.target_size))
        # img_resize = cv2.resize(padded, (self.target_size, self.target_size), interpolation = cv2.INTER_CUBIC)
        # 对ROI进行二值化处理，添加形态学膨胀
        img_resize = padded
        if is_roi:
            img_resize = cv2.resize(padded, (self.target_size, self.target_size), interpolation=cv2.INTER_NEAREST)
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))  # 椭圆核
            img_resize = cv2.dilate(img_resize, kernel, iterations=1) 
            img_resize = (img_resize > 0.5).astype(np.float32)

        # if is_roi:
        #     img_resize = (img_resize > 0.5).astype(np.float32)
    
        return img_resize.astype(np.float32)
    
    def __len__(self):
        return len(self.patient_ids) 

    def __getitem__(self, idx):
        pid = self.patient_ids[idx]
        label = self.labels.loc[pid, 'label']
        
        multimodal_img = []
        multimodal_roi = []
        multimodal_mask = []

        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            img = nib.load(img_path)
            roi = nib.load(roi_path)
            
            img_data = img.get_fdata().astype(np.float32)  # (H, W, D)
            roi_data = roi.get_fdata().astype(np.float32)
            
            # 1. Find slice with largest ROI
            max_roi_area = -1
            max_slice_idx = -1
            for slice_idx in range(roi_data.shape[2]):
                current_area = np.sum(roi_data[..., slice_idx] > 0)
                if current_area > max_roi_area:
                    max_roi_area = current_area
                    max_slice_idx = slice_idx
            
            # Handle no ROI case
            if max_slice_idx == -1:
                print(f"警告：患者 {pid} 的 {mod} 模态没有有效的ROI，使用中间层切片")
                continue
            
            # 2. Calculate center coordinates
            roi_max_slice = roi_data[..., max_slice_idx]
            coords = np.argwhere(roi_max_slice > 0)
            center_y = img_data.shape[0] // 2
            center_x = img_data.shape[1] // 2
            if len(coords) > 0:
                min_y, min_x = coords.min(axis=0)
                max_y, max_x = coords.max(axis=0)
                center_y = (min_y + max_y) // 2
                center_x = (min_x + max_x) // 2
            
            # 3. Get 3 slices (center ±1)
            selected_indices = []
            for offset in range(-1, 2):
                selected_idx = max_slice_idx + offset
                selected_indices.append(selected_idx)
            
            img_slices = []
            roi_slices = []
            mask = []

            for slice_idx in selected_indices:
                # Handle invalid indices
                if slice_idx < 0 or slice_idx >= img_data.shape[2]:
                    print(f"警告：患者 {pid} 的 {mod} 模态没有有效的ROI，使用中间层切片")
                    img_slices.append(np.zeros((self.target_size, self.target_size)))
                    roi_slices.append(np.zeros((self.target_size, self.target_size)))
                    mask.append(0.0)
                    continue
                
                # Load slice data
                img_slice = img_data[..., slice_idx]
                roi_slice = roi_data[..., slice_idx]
                
                # Crop operation
                img_cropped = self._crop_around_center(img_slice, center_y, center_x, is_roi=False)
                roi_cropped = self._crop_around_center(roi_slice, center_y, center_x, is_roi=True)
                
                value = img_cropped.max() - img_cropped.min()
                img_cropped = (img_cropped - img_cropped.min()) /value

                img_cropped = self.transform(img_cropped)

                has_roi = np.any(roi_slice > 0)
                mask.append(1.0 if has_roi else 0.0)
                
                # Data augmentation
                if self.augment:
                    augmented = self.augment(image=standardized, mask=roi_cropped)
                    standardized = augmented['image']
                    roi_cropped = augmented['mask']
                
                img_slices.append(img_cropped) # 保存三张切片
                roi_slices.append(roi_cropped)
            
            multimodal_img.append(torch.FloatTensor(np.stack(img_slices)))   #保存三个方位
            multimodal_roi.append(torch.FloatTensor(np.stack(roi_slices)))
            multimodal_mask.append(torch.FloatTensor(mask))

        multimodal_img = torch.stack(multimodal_img)
        multimodal_roi = torch.stack(multimodal_roi)
        multimodal_mask = torch.stack(multimodal_mask)

        return (
            multimodal_img[0, :],   # (M, D, H, W)，[:,0]代表取每个位面的第一张切片，然后有三个位面
            multimodal_roi[0, :],  # 0代表输入第一个位面的所有切片[0, :]
            # multimodal_mask,
            torch.tensor(label, dtype=torch.long),
            pid
        )

    def apply_augment(self, img, roi):
        """使用torchvision的增强方法"""
        # 转换为Tensor并添加通道维度
        img_tensor = torch.from_numpy(img).unsqueeze(0).float()  # (1,H,W)
        roi_tensor = torch.from_numpy(roi).unsqueeze(0).float()
        
        # 创建同步增强变换
        transform = transforms.Compose([
            transforms.RandomApply([
                transforms.RandomAffine(
                    degrees=10,
                    translate=(0.1, 0.1),
                    scale=(0.9, 1.1),
                    interpolation=transforms.InterpolationMode.BILINEAR
                )
            ], p=0.5),
            transforms.RandomApply([
                transforms.ColorJitter(brightness=0.2, contrast=0.2)
            ], p=0.5)
        ])
        
        # 应用增强
        combined = torch.cat([img_tensor, roi_tensor], dim=0)  # (2,H,W)
        augmented = transform(combined)
        
        # 分离结果并转换回numpy
        img_aug = augmented[0].numpy().squeeze()
        roi_aug = augmented[1].numpy().squeeze()
        
        return img_aug, roi_aug


def visualize_attention(img_batch, roi_batch, labels):
    # 获取首个样本数据
    img, roi, label = img_batch[0].numpy(), roi_batch[0].numpy(), labels[0].item()
    mod_names = ['Axial', 'Sagittal', 'Coronal']
    
    # 创建图形窗口
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    plt.suptitle(f"Label: {label} | Slice: 0/{img.shape[1]-1}", fontsize=16)
    
    # 初始化切片索引
    current_slice = [0]  # 使用列表实现闭包内的可变性
    max_slice = img.shape[1] - 1
    
    # 存储图像对象
    image_plots = [[None for _ in range(3)] for _ in range(3)]
    
    # 初始化显示
    for mod in range(3):
        # 原始图像
        image_plots[0][mod] = axs[0, mod].imshow(img[mod, 0], cmap='gray')
        axs[0, mod].set_title(f"{mod_names[mod]} Image")
        
        # ROI掩膜
        image_plots[1][mod] = axs[1, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.5)
        axs[1, mod].set_title(f"{mod_names[mod]} ROI Mask")
        
        # 叠加显示
        image_plots[2][mod] = axs[2, mod].imshow(img[mod, 0], cmap='gray')
        image_plots[2][mod] = axs[2, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.3)
        axs[2, mod].set_title("Overlay")
    
    def update_display():
        """更新所有模态的显示"""
        fig.suptitle(f"Label: {label} | Slice: {current_slice[0]}/{max_slice}", fontsize=16)
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(img[mod, current_slice[0]])
            
            # 更新ROI掩膜
            image_plots[1][mod].set_data(roi[mod, current_slice[0]])
            
            # 更新叠加显示（需要重新创建以避免透明度混合问题）
            axs[2, mod].images[0].set_data(img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(roi[mod, current_slice[0]])
        
        fig.canvas.draw_idle()
    
    def on_scroll(event):
        """鼠标滚轮事件处理"""
        if event.button == 'up':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.button == 'down':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    # 绑定滚轮事件
    fig.canvas.mpl_connect('scroll_event', on_scroll)
    
    # 添加键盘事件支持
    def on_key(event):
        """键盘事件处理（左右方向键）"""
        if event.key == 'right':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.key == 'left':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # 计算过之后就不需要再多次计算了
    data_transform = transforms.Compose([
                                     transforms.ToTensor(),#PIL Image格式的数据才发挥作用
                                     transforms.Resize((224, 224)),
                                     transforms.Normalize([0.5], [0.5]),
                                     transforms.Lambda(lambda x: x.squeeze(0)),
                                     ])
    
    dataset = MRIDataset(root_dir='D:/PyChrom/PythonProject/2.medical_image/DeepTrip/datasets/pT/three_label_data/train',
                         transform=data_transform)
    
    print(len(dataset))
    print(dataset.class_counts)
    print("有效患者数量:", len(dataset.patient_ids))

    loader = DataLoader(dataset, batch_size=4,shuffle=True)
    
    img_batch, roi_batch, labels, pid = next(iter(loader))

    visualize_attention(img_batch, roi_batch, labels)

    print(f"IMG   shape: {img_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
    print(f"ROI   shape: {roi_batch.shape}")
    # print(f"Batch shape: {mask_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
